#!/usr/bin/env python3
import os
import queue
import sounddevice as sd
import vosk
import sys
import os
import re

class AudioRecognizer:
    def __init__(self):
        self.model_dir = (
            os.path.expanduser("~") + "/Lepi_Data/ros/audio_recognizer/vosk-model-small-cn-0.22"
        )
        self.samplerate = 16000
        self.model = vosk.Model(self.model_dir)
        self.q = queue.Queue()
        devices = sd.query_devices()
        self.usb_device_index = None
        for device in devices:
            if "wm8960" in device["name"]:
                sd.default.device = device["index"]
            if 'MCP01' in device["name"]:
                self.usb_device_index = device["index"]
        if self.usb_device_index is not None:
            sd.default.device = self.usb_device_index
        print(devices)
    def callback(self, indata, frames, time, status):
        """This is called (from a separate thread) for each audio block."""
        if status:
            print(status, file=sys.stderr)
        self.q.put(bytes(indata))

    def recognize(self):
        try:
            with sd.RawInputStream(
                samplerate=self.samplerate,
                blocksize=8000,
                device=None,
                dtype="int16",
                channels=1,
                callback=self.callback,
            ):
                print("#" * 80)
                print("Press Ctrl+C to stop the recording")
                print("#" * 80)

                rec = vosk.KaldiRecognizer(self.model, self.samplerate)
                while True:
                    data = self.q.get()
                    if rec.AcceptWaveform(data):
                        # print(rec.Result())
                        res = re.search( r'"text" : "(.*)"', rec.Result(), re.M|re.I)
                        return res.group(1)
                    else:
                        print(rec.PartialResult())
        except KeyboardInterrupt:
            print("\nDone")
            return ''
        except Exception as e:
            print(type(e).__name__ + ": " + str(e))
            return ''

if __name__ == '__main__':
    ar = AudioRecognizer()
    print(ar.recognize())